-- PostgreSQL 初始化脚本
-- 启用 pgvector 扩展并创建示例表

-- 启用 pgvector 扩展
CREATE EXTENSION IF NOT EXISTS vector;

-- 创建应用用户
CREATE USER appuser WITH PASSWORD 'AppUser2025Secure';

-- 授予权限
GRANT CONNECT ON DATABASE vectordb TO appuser;
GRANT USAGE ON SCHEMA public TO appuser;
GRANT CREATE ON SCHEMA public TO appuser;

-- 创建示例表：文档向量存储
CREATE TABLE IF NOT EXISTS documents (
    id SERIAL PRIMARY KEY,
    title VARCHAR(255) NOT NULL,
    content TEXT NOT NULL,
    embedding vector(1536), -- OpenAI embedding 维度
    metadata JSONB,
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- 创建向量索引（HNSW 算法，适合高维向量相似性搜索）
CREATE INDEX IF NOT EXISTS documents_embedding_idx 
ON documents USING hnsw (embedding vector_cosine_ops);

-- 创建示例表：用户表
CREATE TABLE IF NOT EXISTS users (
    id SERIAL PRIMARY KEY,
    username VARCHAR(50) UNIQUE NOT NULL,
    email VARCHAR(100) UNIQUE NOT NULL,
    password_hash VARCHAR(255) NOT NULL,
    profile_embedding vector(512), -- 用户画像向量
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- 创建用户画像向量索引
CREATE INDEX IF NOT EXISTS users_profile_embedding_idx 
ON users USING hnsw (profile_embedding vector_cosine_ops);

-- 创建示例表：产品表
CREATE TABLE IF NOT EXISTS products (
    id SERIAL PRIMARY KEY,
    name VARCHAR(255) NOT NULL,
    description TEXT,
    category VARCHAR(100),
    price DECIMAL(10,2),
    feature_embedding vector(768), -- 产品特征向量
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- 创建产品特征向量索引
CREATE INDEX IF NOT EXISTS products_feature_embedding_idx 
ON products USING hnsw (feature_embedding vector_cosine_ops);

-- 授予表权限给应用用户
GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO appuser;
GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO appuser;

-- 插入示例数据
INSERT INTO documents (title, content, embedding, metadata) VALUES 
(
    '人工智能简介',
    '人工智能（AI）是计算机科学的一个分支，致力于创建能够执行通常需要人类智能的任务的系统。',
    '[0.1, 0.2, 0.3]'::vector, -- 示例向量（实际应用中应该是真实的嵌入向量）
    '{"category": "technology", "language": "zh", "source": "manual"}'::jsonb
),
(
    'Machine Learning Basics',
    'Machine learning is a subset of artificial intelligence that focuses on algorithms that can learn and make decisions from data.',
    '[0.4, 0.5, 0.6]'::vector,
    '{"category": "technology", "language": "en", "source": "manual"}'::jsonb
);

INSERT INTO users (username, email, password_hash, profile_embedding) VALUES 
(
    'admin',
    'admin@example.com',
    '$2b$12$dummy_hash_for_demo_purposes_only',
    '[0.7, 0.8, 0.9]'::vector
),
(
    'user1',
    'user1@example.com',
    '$2b$12$another_dummy_hash_for_demo_only',
    '[0.1, 0.3, 0.5]'::vector
);

INSERT INTO products (name, description, category, price, feature_embedding) VALUES 
(
    '智能手机',
    '最新款智能手机，配备先进的AI芯片和高清摄像头',
    'electronics',
    999.99,
    '[0.2, 0.4, 0.6]'::vector
),
(
    '机器学习教程',
    '深入浅出的机器学习教程，适合初学者和进阶者',
    'books',
    49.99,
    '[0.8, 0.1, 0.3]'::vector
);

-- 创建向量相似性搜索的辅助函数
CREATE OR REPLACE FUNCTION find_similar_documents(
    query_embedding vector(1536),
    similarity_threshold float DEFAULT 0.7,
    max_results int DEFAULT 10
)
RETURNS TABLE (
    id int,
    title varchar,
    content text,
    similarity float,
    metadata jsonb
) AS $$
BEGIN
    RETURN QUERY
    SELECT 
        d.id,
        d.title,
        d.content,
        1 - (d.embedding <=> query_embedding) as similarity,
        d.metadata
    FROM documents d
    WHERE 1 - (d.embedding <=> query_embedding) > similarity_threshold
    ORDER BY d.embedding <=> query_embedding
    LIMIT max_results;
END;
$$ LANGUAGE plpgsql;

-- 创建更新时间戳的触发器函数
CREATE OR REPLACE FUNCTION update_updated_at_column()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = CURRENT_TIMESTAMP;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- 为所有表创建更新时间戳触发器
CREATE TRIGGER update_documents_updated_at BEFORE UPDATE ON documents
    FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();

CREATE TRIGGER update_users_updated_at BEFORE UPDATE ON users
    FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();

CREATE TRIGGER update_products_updated_at BEFORE UPDATE ON products
    FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();

-- 输出初始化完成信息
DO $$
BEGIN
    RAISE NOTICE 'PostgreSQL with pgvector 初始化完成！';
    RAISE NOTICE '- 启用了 pgvector 扩展';
    RAISE NOTICE '- 创建了应用用户: appuser';
    RAISE NOTICE '- 创建了向量存储表和索引';
    RAISE NOTICE '- 插入了示例数据';
    RAISE NOTICE '- 创建了向量相似性搜索函数';
END $$;
